#!/usr/bin/env python3
import numpy as np
import open3d as o3d
import cv2
import matplotlib.pyplot as plt
from matplotlib.path import Path

def load_point_cloud(ply_path, z_min=0.1, z_max=6.0):
    """读取 .ply 点云文件并应用深度过滤"""
    pcd = o3d.io.read_point_cloud(ply_path)
    if not pcd.has_points():
        raise ValueError("点云文件为空")
    
    # 提取点和颜色
    points = np.asarray(pcd.points)
    colors = np.asarray(pcd.colors) if pcd.has_colors() else None
    
    # 过滤无效点
    valid_mask = (points[:, 2] > z_min) & (points[:, 2] < z_max) & \
                 (~np.isnan(points[:, 0])) & (~np.isnan(points[:, 1])) & (~np.isnan(points[:, 2]))
    points = points[valid_mask]
    
    if len(points) == 0:
        raise ValueError("所有点无效，请检查深度范围")
    
    # 创建新的点云对象
    filtered_pcd = o3d.geometry.PointCloud()
    filtered_pcd.points = o3d.utility.Vector3dVector(points)
    
    # 应用颜色（如果存在）
    if colors is not None and len(colors[valid_mask]) == len(points):
        filtered_pcd.colors = o3d.utility.Vector3dVector(colors[valid_mask])
    
    print(f"加载点云，原始点数: {len(pcd.points)}, 过滤后点数: {len(points)}")
    print(f"点云坐标范围: x=[{points[:, 0].min():.2f}, {points[:, 0].max():.2f}], "
          f"y=[{points[:, 1].min():.2f}, {points[:, 1].max():.2f}], "
          f"z=[{points[:, 2].min():.2f}, {points[:, 2].max():.2f}]")
    return filtered_pcd, valid_mask

def project_points_to_image(points, fx, fy, cx, cy, image_width=1280, image_height=720):
    """将 3D 点云投影到 2D 图像平面，过滤无效点"""
    points_np = np.asarray(points)
    # 过滤 z <= 0 或 nan 的点
    valid_mask = (points_np[:, 2] > 0) & (~np.isnan(points_np[:, 2]))
    points_np = points_np[valid_mask]
    
    # 投影到 2D
    u = (points_np[:, 0] * fx / points_np[:, 2] + cx)
    v = (points_np[:, 1] * fy / points_np[:, 2] + cy)
    
    # 转换为整数并限制在图像范围内
    u = np.clip(u, 0, image_width - 1).astype(int)
    v = np.clip(v, 0, image_height - 1).astype(int)
    
    # 打印投影点范围
    print(f"有效投影点数: {len(points_np)}")
    print(f"投影点范围: u=[{u.min():.2f}, {u.max():.2f}], v=[{v.min():.2f}, {v.max():.2f}]")
    return np.stack((u, v), axis=1), valid_mask

def is_point_in_polygon_boundary(points_2d, polygon, boundary_thickness_pixels=10):
    """检查 2D 点是否在多边形边界附近"""
    path = Path(polygon)
    # 计算点到多边形边界的距离
    boundary_mask = np.zeros(len(points_2d), dtype=bool)
    for i, point in enumerate(points_2d):
        dist = path.contains_point(point, radius=boundary_thickness_pixels)
        if dist != 0:  # 非零表示在边界附近
            boundary_mask[i] = True
    print(f"边界点数 (2D): {np.sum(boundary_mask)}")
    return boundary_mask

def extract_boundary_points(ply_path, polygon, fx, fy, cx, cy, image_width=1280, image_height=720, boundary_thickness_pixels=10, z_min=0.1, z_max=6.0):
    """从点云中提取多边形边界上的点"""
    # 加载点云并应用深度过滤
    pcd, valid_mask = load_point_cloud(ply_path, z_min, z_max)
    points = np.asarray(pcd.points)
    
    # 投影到 2D 图像平面
    points_2d, proj_valid_mask = project_points_to_image(points, fx, fy, cx, cy, image_width, image_height)
    
    # 筛选在多边形边界上的点
    boundary_mask_2d = is_point_in_polygon_boundary(points_2d, polygon, boundary_thickness_pixels)
    
    # 映射回 3D 点云
    boundary_mask = np.zeros(len(points), dtype=bool)
    boundary_mask[proj_valid_mask] = boundary_mask_2d
    
    # 提取边界点
    boundary_points = points[boundary_mask]
    if len(boundary_points) == 0:
        raise ValueError("未提取到边界点，请检查多边形坐标、边界厚度或相机内参")
    
    boundary_pcd = o3d.geometry.PointCloud()
    boundary_pcd.points = o3d.utility.Vector3dVector(boundary_points)
    
    # 如果有点云颜色，提取对应颜色
    if pcd.has_colors():
        colors = np.asarray(pcd.colors)
        boundary_pcd.colors = o3d.utility.Vector3dVector(colors[boundary_mask])
    
    print(f"提取到边界点数 (3D): {len(boundary_points)}")
    return boundary_pcd, points_2d, boundary_mask_2d

def main():
    # 参数设置
    ply_path = "/root/ros_ws/src/OrbbecSDK_ROS1/scripts/data/pointclouds/orbbec_20250725_173354.ply"  # 替换为实际 .ply 文件路径
    # 相机内参（替换为实际值）
    fx, fy = 525.0, 525.0  # 焦距
    cx, cy = 640.0, 360.0  # 主点（假设 1280x720 分辨率）
    image_width, image_height = 1280, 720  # 图像分辨率
    z_min, z_max = 0.1, 6.0  # 深度范围
    
    # 定义更大的多边形，覆盖图像中心区域
    polygon = np.array([
        [400, 200],
        [880, 200],
        [880, 520],
        [400, 520]
    ])  # 调整为更大的矩形，靠近图像中心
    
    # 提取边界点
    try:
        boundary_pcd, points_2d, boundary_mask_2d = extract_boundary_points(
            ply_path, polygon, fx, fy, cx, cy, image_width, image_height, boundary_thickness_pixels=10, z_min=z_min, z_max=z_max
        )
        
        # 可视化 3D 点云
        o3d.visualization.draw_geometries([boundary_pcd], window_name="Boundary Points", point_show_normal=False, point_size=3.0)
        
        # 可视化 2D 多边形和投影点
        plt.figure()
        plt.plot(polygon[:, 0], polygon[:, 1], 'b-', label="Polygon")
        plt.plot(polygon[[0, -1], 0], polygon[[0, -1], 1], 'b-')  # 闭合多边形
        plt.scatter(points_2d[:, 0], points_2d[:, 1], c='gray', s=1, alpha=0.1, label="All Points")
        plt.scatter(points_2d[boundary_mask_2d, 0], points_2d[boundary_mask_2d, 1], c='r', s=5, label="Boundary Points")
        plt.legend()
        plt.title("2D Projection of Points")
        plt.xlabel("u (pixels)")
        plt.ylabel("v (pixels)")
        plt.xlim(0, image_width)
        plt.ylim(0, image_height)
        plt.gca().invert_yaxis()  # 图像坐标系 Y 轴通常向下
        plt.show()
    
    except ValueError as e:
        print(f"错误: {str(e)}")

if __name__ == "__main__":
    main()
